#!/usr/bin/env python3

import os, sys
import numpy as np
import scipy as sp
import matplotlib.pylab as plt
from tqdm import trange

sys.path.append('./python')

#from dcds_mask_44S import *
from dcds_mask import *
from getOptWeightsOneFrame import *
from getOptWeightsTwoFrames import *

gain = 1

res_dir = 'weights'
if os.path.isdir(res_dir) is False:
    os.system('mkdir '+res_dir)


def Run_Opt(n_pair=4, file_num=180, nsample_per_pixel=108, gain=1.0):
    
#    mask = DCDS_Mask_44S()
    mask = DCDS_Mask()
    mask = mask.gen_mask(n_pair)

    n_ref = np.sum(mask==1)
    n_sig = np.sum(mask==-1)
    print('> n_ref = {}, n_sig = {}'.format(n_ref, n_sig))
    
    std_old = []
    std_new = []

    fig = plt.figure(figsize=(16,10))

    for i in trange(16):
        print('> processing ch: {}'.format(i+1))
        os_wave_dir = 'OS_'+str(i+1)
        wfile = res_dir + '/' + 'w_os_' + str(n_pair) + '_' + str(i+1)+'.txt'
        std_oldx, std_newx = getWeightsOneFrame(os_wave_dir, file_num=file_num, \
                                                mask=mask, nsample_per_pixel=nsample_per_pixel, \
                                                n_ref=n_ref, n_sig=n_sig, out_fname=wfile)
        std_old.append(std_oldx)
        std_new.append(std_newx)
        tmp = np.loadtxt(wfile)
        plt.subplot(4,4,i+1)
        plt.plot(tmp,'--.')
        plt.title('ch :'+str(i+1))
    
    plt.tight_layout()
    plt.savefig('w_'+str(n_pair)+'.png')

    plt.figure()

    std_old = np.array(std_old)
    std_new = np.array(std_new)

    plt.plot(std_old*gain, 'r:o', label='raw_'+str(n_pair))
    plt.plot(std_new*gain, 'g--o', label='opt_'+str(n_pair))
    plt.grid(linestyle='dashed')
    
    plt.ylabel('e-')
    ymin = min( min(std_old*gain), min(std_new*gain) ) * 0.9
    ymax = max( max(std_old*gain), max(std_new*gain) ) * 1.1
    plt.ylim([ymin,ymax])

    plt.legend(ncol=2)
    plt.savefig('RMS_1fs_'+str(n_pair)+'.png')
    np.savetxt('RMS_1fs_'+str(n_pair)+'.txt', np.vstack((std_old,std_new)).transpose(), fmt='%8.3f');
    # plt.show()

###################################################################################################
if __name__ == '__main__':
    file_num = 200
    Run_Opt(n_pair=6, file_num=file_num, nsample_per_pixel=20)
